import os.path

import torch.nn
from torch.utils.data import DataLoader
from tqdm import tqdm

from MySAM2 import SAM2
from MyData import MyDataset
from tensorboardX import SummaryWriter
from torch.optim import lr_scheduler
from EvaluatorModule import Evaluator
import random
import numpy as np


def test(model, dataset, evaluator):
    with torch.no_grad():
        # 设置为eval模式
        model.sam2_model.sam_prompt_encoder.eval()
        model.sam2_model.sam_mask_decoder.eval()
        status = {'f1': 0, 'rec': 0, 'pre': 0, 'qua': 0, 'com': 0, 'cor': 0}
        with tqdm(total=len(dataset), desc='{:<5} {:<5} {:<3}'.format('test', 'epoch', 0)) as pbar:
            for idx, data in enumerate(dataset):
                # 正向传播
                out = model.forward(image=data['img'])[0]
                # 计算评估指标
                pre, rec, f1 = evaluator.evaluateSingleFR(torch.where(out >= 0.5, 1.0, 0), data['mask'])
                cor, com, qua = evaluator.evaluateSingleTP(torch.where(out >= 0.5, 1.0, 0), data['mask'])
                # 累积评估指标
                status['pre'] += pre
                status['rec'] += rec
                status['f1'] += f1
                status['com'] += com
                status['cor'] += cor
                status['qua'] += qua
                info = {'Pre': status['pre'] / (idx + 1),
                        'Rec': status['rec'] / (idx + 1),
                        'F1': status['f1'] / (idx + 1),
                        'Com': status['com'] / (idx + 1),
                        'Cor': status['cor'] / (idx + 1),
                        'Qua': status['qua'] / (idx + 1)
                        }
                # 展示信息
                pbar.set_postfix(info)
                pbar.update()
        print({'Pre': status['pre'] / (idx + 1),
               'Rec': status['rec'] / (idx + 1),
               'F1': status['f1'] / (idx + 1),
               'Com': status['com'] / (idx + 1),
               'Cor': status['cor'] / (idx + 1),
               'Qua': status['qua'] / (idx + 1)
               })


if __name__ == '__main__':
    # 数据集
    dataset = MyDataset(r'Datasets/Crack500', 'test', 'cuda')
    # 模型
    model = SAM2(config_file=r'configs/sam2.1_hiera_l.yaml', ckpt_path='weights/sam2.1_hiera_large.pt')
    #model.load_state_dict(torch.load("checkpoints/SAM2FPN-Mask/sam_finetuned_bestQua.pth"))
    # 评估器
    evaluator = Evaluator(benchmark='Crack500')
    # 最佳指标
    test(model, dataset, evaluator)
